import numpy as np  
import datetime
import itertools
import os.path as osp
import tempfile
from collections import OrderedDict
from typing import Dict, List, Optional, Sequence, Union

import numpy as np
from scipy.spatial.distance import cdist   
  

import torch
from mmengine.evaluator import BaseMetric
from mmengine.fileio import dump, get_local_path, load
from mmengine.logging import MMLogger
from terminaltables import AsciiTable

from mmdet.datasets.api_wrappers import COCO, COCOeval
from mmdet.registry import METRICS
from mmdet.structures.mask import encode_mask_results
from ..functional import eval_recalls
from concurrent.futures import ThreadPoolExecutor
from tqdm import tqdm as tqdm
import json

@METRICS.register_module()
class Slot_Entrance_Metric(BaseMetric):
    def __init__(self, 
                 threshold = 20,
                 collect_device = 'cpu', 
                 prefix = None, 
                 collect_dir = None) -> None:
        super().__init__(collect_device=collect_device, prefix=prefix)
        self.threshold = threshold

    def match_points(self,A, B):  
        """  
        匹配点集A和B，根据距离进行匹配，阈值设为threshold。  
        """  
        # 计算点集A和B之间的距离矩阵  
        if len(A) == 0 or len(B) == 0:  
            return -1, -1 
        dist_matrix = cdist(A, B)  
        
        # 找到点集A中每个点距离最近的点在点集B中的索引  
        indices = np.argmin(dist_matrix, axis=1)  
        
        # 根据索引找到点集B中每个点的匹配点在点集A中的索引  
        matches = np.argmin(dist_matrix, axis=0)[indices]  
        
        # 检查匹配是否满足阈值要求  
        valid_matches = dist_matrix[np.arange(len(A)), indices] < self.threshold  
        matches = matches[valid_matches]  
        indices = indices[valid_matches]  
        
        # 返回匹配结果  
        return matches, indices

    def cal_onebatch(self,pred,gt):
        slot_pred_match_id,slot_gt_match_id = self.match_points(pred["center"][0].reshape(-1,2),gt[0][:,0:2].reshape(-1,2))

        slot_pred_num = pred["center"][0].shape[0]
        slot_gt_num = gt[0].shape[0]
        if isinstance(slot_pred_match_id,int) and isinstance(slot_gt_match_id,int):
            slot_tp_num = 0
            slot_center_error = np.empty([0,2])
            slot_joint_left_error =  np.empty([0,2])
            slot_joint_right_error = np.empty([0,2])
            slot_ori_error = np.empty([0,1])
        else:
            slot_tp_num = slot_pred_match_id.shape[0]
            slot_center_error = pred["center"][0][slot_pred_match_id] - gt[0][slot_gt_match_id][:,0:2]
            slot_joint_left_error = pred["joint_l"][0][slot_pred_match_id] - gt[0][slot_gt_match_id][:,2:4]
            slot_joint_right_error = pred["joint_r"][0][slot_pred_match_id] - gt[0][slot_gt_match_id][:,4:6]
            pred_ori = np.arctan2(pred["ori"][0][:,1],pred["ori"][0][:,0])
            pred_ori[pred_ori<0] = 2*np.pi
            slot_ori_error = pred_ori[slot_pred_match_id].reshape(-1,1) - gt[0][slot_gt_match_id][:,6:7]




        joint_pred_match_id,joint_gt_match_id = self.match_points(pred['joint_center'][0].reshape(-1,2),gt[1].reshape(-1,2))
        joint_pred_num = pred["joint_center"][0].shape[0]
        joint_gt_num = gt[1].shape[0]
        if isinstance(joint_pred_match_id,int) and isinstance(joint_pred_match_id,int):
            joint_tp_num = 0
            joint_center_error = np.empty([0,2])
        else:
            joint_tp_num = joint_pred_match_id.shape[0]
            joint_center_error = pred["joint_center"][0][joint_pred_match_id] - gt[1][joint_gt_match_id]


        return [np.array([slot_tp_num,slot_pred_num,slot_gt_num,joint_tp_num,joint_pred_num,joint_gt_num]).reshape(1,6),\
            np.concatenate([slot_center_error,slot_joint_left_error,slot_joint_right_error,slot_ori_error],axis=1),\
            joint_center_error]


    def process(self,data_batch,data_samples):
#data_samples (Sequence[BaseDataElement]): predictions of the model,
#     and the ground truth of the validation set.
# data_batch (Any, optional): A batch of data from the dataloader.
        bs = len(data_batch['inputs'])
        for i in range(bs):
            gt = [data_samples[i]['slots'],data_samples[i]['mark_points']]
            pred = data_samples[i]['pred_instances']
            self.results.append((gt,pred))
    def compute_metrics(self, results: list) -> dict:
        logger: MMLogger = MMLogger.get_current_instance()
        logger.info("match done,compute metrics")
        gts, preds = zip(*results)

        # for i in range(len(gts)):
        #     result = self.cal_onebatch(preds[i],gts[i])

        with ThreadPoolExecutor(max_workers=32) as pool:
            threads = [pool.submit(self.cal_onebatch,preds[i],gts[i]) for i in range(len(self.results))]
            # threads = tqdm(threads)
            cal_results = [t.result() for t in threads]

        num_results = [result[0] for result in cal_results]
        slot_error_results = [result[1] for result in cal_results]
        joint_error_results = [result[2] for result in cal_results]

        num_results = np.concatenate(num_results,axis=0)
        slot_error_results = np.concatenate(slot_error_results,axis=0)
        joint_error_results = np.concatenate(joint_error_results,axis=0)

        slot_precision = num_results[:,0].sum() / num_results[:,1].sum()
        slot_recall = num_results[:,0].sum() / num_results[:,2].sum()

        joint_precision = num_results[:,3].sum() / num_results[:,4].sum()
        joint_recall = num_results[:,3].sum() / num_results[:,5].sum()

        slot_entrance_x_error = np.abs(slot_error_results[:,0]).mean()
        slot_entrance_y_error = np.abs(slot_error_results[:,1]).mean()
        slot_joint_l_x_error = np.abs(slot_error_results[:,2]).mean()
        slot_joint_l_y_error = np.abs(slot_error_results[:,3]).mean()
        slot_joint_r_x_error = np.abs(slot_error_results[:,4]).mean()
        slot_joint_r_y_error = np.abs(slot_error_results[:,5]).mean()
        slot_ori_error = np.abs(slot_error_results[:,6]).mean()

        joint_x_error = np.abs(joint_error_results[:,0]).mean()
        joint_y_error = np.abs(joint_error_results[:,1]).mean()

        eval_results = dict(
                            slot_precision = slot_precision,
                            slot_recall = slot_recall,
                            joint_precision = joint_precision,
                            joint_recall = joint_recall,
                            slot_entrance_x_error = slot_entrance_x_error,
                            slot_entrance_y_error=slot_entrance_y_error,
                            slot_joint_l_x_error=slot_joint_l_x_error,
                            slot_joint_l_y_error=slot_joint_l_y_error,
                            slot_joint_r_x_error=slot_joint_r_x_error,
                            slot_joint_r_y_error=slot_joint_r_y_error,
                            slot_ori_error=slot_ori_error,
                            joint_x_error=joint_x_error,
                            joint_y_error=joint_y_error)
        return eval_results

